package study.flink;

import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.functions.FlatMapFunction;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.SinkFunction;
import org.apache.flink.streaming.api.functions.source.SocketTextStreamFunction;
import org.apache.flink.streaming.api.functions.windowing.RichWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;

import java.time.Duration;
import java.util.Date;

public class WordCount {

    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);
        env.addSource(new SocketTextStreamFunction("vm-ubuntu", 8888, "\n", 1))
                .flatMap(new FlatMapFunction<String, String>() {
                    @Override
                    public void flatMap(String value, Collector<String> out) throws Exception {
                        for (char c : value.toCharArray()) {
                            out.collect(String.valueOf(c));
                        }
                    }
                })
                .filter(value -> !"".equals(value))
                .map(new MapFunction<String, Tuple2<String, Long>>() {
                    @Override
                    public Tuple2<String, Long> map(String value) throws Exception {
                        Tuple2<String, Long> r = new Tuple2<>(value, System.currentTimeMillis());
                        System.out.println(new Date((Long) r.getField(1)));
                        return r;

                    }
                })
                .assignTimestampsAndWatermarks(
                        WatermarkStrategy.<Tuple2<String, Long>>forBoundedOutOfOrderness(Duration.ofSeconds(1))
                                .withTimestampAssigner((element, recordTimestamp) -> element.getField(1)))
                .keyBy(new KeySelector<Tuple2<String, Long>, String>() {
                    @Override
                    public String getKey(Tuple2<String, Long> value) throws Exception {
                        return value.getField(0);
                    }
                })
                .window(TumblingEventTimeWindows.of(Time.seconds(3)))
                .aggregate(new AggregateFunction<Tuple2<String, Long>, Integer, Integer>() {
                    @Override
                    public Integer createAccumulator() {
                        return 0;
                    }

                    @Override
                    public Integer add(Tuple2<String, Long> value, Integer accumulator) {
                        return accumulator + 1;
                    }

                    @Override
                    public Integer getResult(Integer accumulator) {
                        return accumulator;
                    }

                    @Override
                    public Integer merge(Integer a, Integer b) {
                        return a + b;
                    }
                }, new RichWindowFunction<Integer, Tuple3<String, Integer, Long>, String, TimeWindow>() {
                    ValueState<Integer> total;

                    @Override
                    public void open(Configuration parameters) throws Exception {
                        total = getRuntimeContext().getState(new ValueStateDescriptor<Integer>("total", Integer.class));
                    }

                    @Override
                    public void close() throws Exception {
                        total.clear();
                    }

                    @Override
                    public void apply(String s, TimeWindow window, Iterable<Integer> input, Collector<Tuple3<String, Integer, Long>> out) throws Exception {
                        Integer total = this.total.value();
                        if (total == null) {
                            total = 0;
                        }
                        int count = input.iterator().next();
                        total = total + count;
                        this.total.update(total);
                        Tuple3<String, Integer, Long> res = new Tuple3<>(s, total, window.getEnd());
                        System.out.println(String.format("windowEnd: %s", window.getEnd()));
                        out.collect(res);
                    }
                })
                .addSink(new SinkFunction<Tuple3<String, Integer, Long>>() {

                    @Override
                    public void invoke(Tuple3<String, Integer, Long> value, Context context) throws Exception {

                        System.out.printf("now：%s, windowEnd=%s, key=%s, num=%s%n", new Date(), new Date((Long) value.getField(2)), value.getField(0), value.getField(1));

                    }
                });

        env.execute();

    }

}
